use bytemuck::{Pod, Zeroable};
use embedded_graphics::{
    pixelcolor::{Gray8, Rgb888},
    prelude::{GrayColor, RgbColor},
};
use std::{
    fs,
    io::{self, Write},
    ops::{Index, IndexMut},
    path::Path,
};
use tinytga::{Bpp, RawTga, Tga};

#[repr(C, packed)]
#[derive(Pod, Zeroable, Copy, Clone, Default, Debug)]
pub struct TGAHeader {
    idlength: u8,
    colormaptype: u8,
    datatypecode: u8,
    colormaporigin: u16,
    colormaplength: u16,
    colormapdepth: u8,
    x_origin: u16,
    y_origin: u16,
    width: u16,
    height: u16,
    bitsperpixel: u8,
    imagedescriptor: u8,
}

#[derive(Copy, Clone, Default, Debug)]
pub struct TGAColor {
    bgra: [u8; 4],
}

impl TGAColor {
    pub const fn from_bgra(bgra: &[u8; 4]) -> Self {
        TGAColor { bgra: *bgra }
    }
}

impl Index<usize> for TGAColor {
    type Output = u8;
    fn index(&self, index: usize) -> &Self::Output {
        &self.bgra[index]
    }
}

impl IndexMut<usize> for TGAColor {
    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
        &mut self.bgra[index]
    }
}

#[derive(Clone, Copy)]
pub enum ColorFormat {
    GRAY,
    RGBA,
}

impl ColorFormat {
    fn bpp_value(&self) -> u8 {
        match self {
            ColorFormat::GRAY => 1,
            ColorFormat::RGBA => 4,
        }
    }
}

#[derive(Default, Clone)]
pub struct TGAImage {
    w: usize,
    h: usize,
    bpp: u8,
    pub data: Vec<u8>,
}

impl TGAImage {
    pub fn new(w: usize, h: usize, format: ColorFormat) -> Self {
        TGAImage {
            w,
            h,
            bpp: format.bpp_value(),
            data: vec![0; w * h * format.bpp_value() as usize],
        }
    }

    pub fn set_data(mut self, data: &[u8]) -> Self {
        self.data.copy_from_slice(data);
        self
    }

    pub fn width(&self) -> usize {
        self.w
    }

    pub fn height(&self) -> usize {
        self.h
    }

    pub fn from_tga_file<P: AsRef<Path>>(filename: P) -> io::Result<Self> {
        assert!(filename.as_ref().exists());

        let data = std::fs::read(filename.as_ref()).unwrap();
        let raw_tga = RawTga::from_slice(&data).unwrap();
        let header = raw_tga.header();

        match header.pixel_depth {
            Bpp::Bits24 => {
                let w = header.width as usize;
                let h = header.height as usize;
                let img: Tga<Rgb888> = Tga::from_slice(&data).unwrap();
                let data = img
                    .pixels()
                    .map(|x| {
                        let mut v = vec![0u8; 4];
                        v[0] = x.1.r();
                        v[1] = x.1.g();
                        v[2] = x.1.b();
                        v[3] = 255;
                        v
                    })
                    .collect::<Vec<_>>()
                    .concat();
                assert!(w * h * 4 == data.len());
                Ok(TGAImage {
                    w,
                    h,
                    bpp: ColorFormat::RGBA.bpp_value(),
                    data,
                })
            }
            Bpp::Bits32 => {
                let img = image::open(filename).unwrap();
                let rgba_img = img.to_rgba8();
                let w = rgba_img.width() as usize;
                let h = rgba_img.height() as usize;
                let data = rgba_img.into_vec();
                assert!(w * h * 4 == data.len());

                Ok(TGAImage {
                    w,
                    h,
                    bpp: ColorFormat::RGBA.bpp_value(),
                    data,
                })
            }
            Bpp::Bits8 => {
                let img: Tga<Gray8> = Tga::from_slice(&data).unwrap();
                let data = img.pixels().map(|x| x.1.luma()).collect::<Vec<_>>();
                let w = header.width as usize;
                let h = header.height as usize;
                assert!(w * h == data.len());
                Ok(TGAImage {
                    w,
                    h,
                    bpp: ColorFormat::GRAY.bpp_value(),
                    data,
                })
            }
            _ => {
                unimplemented!(
                    "{:?} format {:?} not supported",
                    filename.as_ref(),
                    header.pixel_depth
                )
            }
        }
    }

    pub fn write_tga_file<P: AsRef<Path>>(&self, filename: P, vflip: bool) -> io::Result<()> {
        assert!(self.bpp == ColorFormat::RGBA.bpp_value());

        let developer_area_ref = [0u8; 4];
        let extension_area_ref = [0u8; 4];
        let footer: [u8; 18] = [
            b'T', b'R', b'U', b'E', b'V', b'I', b'S', b'I', b'O', b'N', b'-', b'X', b'F', b'I',
            b'L', b'E', b'.', b'\0',
        ];

        let mut f = fs::File::create(filename)?;

        let header = TGAHeader {
            datatypecode: 2,
            width: self.w as u16,
            height: self.h as u16,
            bitsperpixel: self.bpp << 3,
            imagedescriptor: if vflip { 0x00 } else { 0x20 },
            ..Default::default()
        };

        f.write_all(bytemuck::bytes_of(&header))?;

        let data = (0..self.w * self.h)
            .map(|i| {
                let v = self.data[i * 4..(i + 1) * 4].to_vec();
                vec![v[2], v[1], v[0], v[3]]
            })
            .collect::<Vec<_>>()
            .concat();
        f.write_all(&data)?;

        f.write_all(&developer_area_ref)?;
        f.write_all(&extension_area_ref)?;
        f.write_all(&footer)?;

        Ok(())
    }
}
